from django.db.models.query import QuerySet, ValuesQuerySet
from django.http import HttpResponse

####################################
## from http://www.djangosnippets.org/snippets/1151/ by Tarken
## CSVResponse is original, though adapted
####################################

class ExcelResponse(HttpResponse):
    
    def __init__(self, data, output_name='excel_data', headers=None, force_csv=False):
        
        # Make sure we've got the right type of data to work with
        valid_data = False
        if isinstance(data, ValuesQuerySet):
            data = list(data)
        elif isinstance(data, QuerySet):
            data = list(data.values())
        if hasattr(data, '__getitem__'):
            if isinstance(data[0], dict):
                if headers is None:
                    headers = data[0].keys()
                data = [[row[col] for col in headers] for row in data]
                data.insert(0, headers)
            if hasattr(data[0], '__getitem__'):
                valid_data = True
        assert valid_data is True, "ExcelResponse requires a sequence of sequences"
        
        import cStringIO
        output = cStringIO.StringIO()
        # Excel has a limit on number of rows; if we have more than that, make a csv
        use_xls = False
        if len(data) <= 65535 and force_csv is not True:
            try:
                import xlwt
            except ImportError:
                # xlwt doesn't exist; fall back to csv
                pass
            else:
                use_xls = True
        if use_xls:
            book = xlwt.Workbook()
            sheet = book.add_sheet('Sheet 1')
            for rowx, row in enumerate(data):
                for colx, value in enumerate(row):
                    sheet.write(rowx, colx, str(value))
            book.save(output)
            mimetype = 'application/vnd.ms-excel'
            file_ext = 'xls'
        else:
            import csv
            
            writer = csv.writer(output, dialect='excel')
            writer.writerows(data)
            mimetype = 'text/csv'
            file_ext = 'csv'
        output.seek(0)
        super(ExcelResponse, self).__init__(content=output.getvalue(),
                                            mimetype=mimetype)
        self['Content-Disposition'] = 'attachment;filename="%s.%s"' % \
            (output_name.replace('"', '\"'), file_ext)
            
class CSVResponse(HttpResponse):
    
    def __init__(self, data, output_name='csv_data', headers=None):
        
        # Make sure we've got the right type of data to work with
        valid_data = False
        if isinstance(data, ValuesQuerySet):
            data = list(data)
        elif isinstance(data, QuerySet):
            data = list(data.values())
        if hasattr(data, '__getitem__'):
            if isinstance(data[0], dict):
                if headers is None:
                    headers = data[0].keys()
                data = [[row[col] for col in headers] for row in data]
                data.insert(0, headers)
            if hasattr(data[0], '__getitem__'):
                valid_data = True
        assert valid_data is True, "CSVResponse requires a sequence of sequences"
        
        import cStringIO
        output = cStringIO.StringIO()

        import csv
        
        writer = csv.writer(output, dialect='excel')
        writer.writerows(data)
        mimetype = 'text/csv'
        file_ext = 'csv'
        output.seek(0)
        super(CSVResponse, self).__init__(content=output.getvalue(),
                                            mimetype=mimetype)
        self['Content-Disposition'] = 'attachment;filename="%s.%s"' % \
            (output_name.replace('"', '\"'), file_ext)